import base64
import json
import os
import logging
import ssl
import sys
import time
from collections import OrderedDict
from datetime import datetime, timezone

from aiohttp import web, BasicAuth, ClientSession
import asyncio

import requests

import mh_config
from register import apple_cryptography, pypush_gsa_icloud

logger = logging.getLogger()
logging.basicConfig(level=logging.DEBUG)

def authenticate(self):
    endpoint_user = mh_config.getEndpointUser()
    endpoint_pass = mh_config.getEndpointPass()
    if (endpoint_user is None or endpoint_user == "") and (endpoint_pass is None or endpoint_pass == ""):
        return True
    auth_header = self.headers.get('authorization')
    if auth_header:
        auth_type, auth_encoded = auth_header.split(None, 1)
        if auth_type.lower() == 'basic':
            auth_decoded = base64.b64decode(auth_encoded).decode('utf-8')
            username, password = auth_decoded.split(':', 1)
            if username == endpoint_user and password == endpoint_pass:
                return True
    return False

def getAuth(regenerate=False, second_factor='sms'):
    if os.path.exists(mh_config.getConfigFile()) and not regenerate:
        with open(mh_config.getConfigFile(), "r") as f:
            j = json.load(f)
    else:
        mobileme = pypush_gsa_icloud.icloud_login_mobileme(username=mh_config.USER, password=mh_config.PASS)
        logger.debug('Mobileme result: ' + mobileme)
        j = {'dsid': mobileme['dsid'], 'searchPartyToken': mobileme['delegates']
             ['com.apple.mobileme']['service-data']['tokens']['searchPartyToken']}
        with open(mh_config.getConfigFile(), "w") as f:
            json.dump(j, f)
    return BasicAuth(j['dsid'], j['searchPartyToken'])

async def handle_options(request):
    headers = {
        'Access-Control-Allow-Origin': '*',
        'Access-Control-Allow-Methods': 'GET, OPTIONS',
        'Access-Control-Allow-Headers': 'X-Requested-With, Content-Type, Authorization',
        'Access-Control-Allow-Private-Network': 'true',
    }
    return web.Response(status=200, headers=headers)

async def handle_get(request):
    if not authenticate(request):
        headers = {
            'WWW-Authenticate': 'Basic realm="Auth Realm"'
        }
        return web.Response(status=401, headers=headers)
    headers = {
        'Access-Control-Allow-Origin': '*',
        'Content-type': 'text/plain'
    }
    return web.Response(text="Nothing to see here.", headers=headers)


async def fetch(session, sem, url, data, auth, headers):
    async with sem:  # 限制并发数
        async with session.post(url,
                                auth=auth,
                                headers=headers,
                                json=data) as resp:
            #resp.raise_for_status()
            status = resp.status
            if status != 401:
                resp.raise_for_status()
            else:
                return status, list()
            results = (await resp.json())['results']
            #print(f'{status}: {len(results)} reports received.')
            return status, results
        
async def handle_post(request):
    if not authenticate(request):
        headers = {
            'WWW-Authenticate': 'Basic realm="Auth Realm"'
        }
        return web.Response(status=401, headers=headers)
    try:
        # 解析请求体json
        body = await request.json()
        days = body.get('days', 7)
        #logger.debug(f'Querying for {days} days')

        unixEpoch = int(datetime.now().timestamp())
        startdate = unixEpoch - (60 * 60 * 24 * days)

        ids = list(body['ids'])
        # Date is always 1, because it has no effect
        data = {"search": [{"startDate": 1, "ids": ids}]}

        times_401=0

        hashed_adv_groups = [ids[i:i+3] for i in range(0, len(ids), 3)] # 3个公钥一组

        auth = getAuth(regenerate=False, second_factor='sms')
        anisette_header_time = time.time()
        headers = pypush_gsa_icloud.generate_anisette_headers()

        url = "https://gateway.icloud.com/acsnservice/fetch"

        sem = request.app['sem']

        async with ClientSession() as session:
            tasks = []
            for hashed_adv_group in hashed_adv_groups:
                data = {"search": [{"startDate": startdate * 1000, "endDate": unixEpoch * 1000, "ids": hashed_adv_group}]}

                if time.time() - anisette_header_time > 59:
                    anisette_header_time = time.time()
                    headers = pypush_gsa_icloud.generate_anisette_headers()

                task = asyncio.create_task(fetch(session, sem, url, data, auth, headers))
                tasks.append(task)

            responses = await asyncio.gather(*tasks, return_exceptions=True)
            newResults = OrderedDict()

            for tup in responses:
                if isinstance(tup, Exception):
                    print(f"Request failed: {tup}")
                    continue
                status = tup[0]
                if status == 401:
                    times_401 += 1
                else:
                    res = tup[1]
                    for entry in res:
                        data_bytes = base64.b64decode(entry['payload'])
                        timestamp = int.from_bytes(data_bytes[0:4], 'big') + 978307200
                        if timestamp > startdate:
                            newResults[timestamp] = entry

        if times_401 == len(hashed_adv_groups):
            raise RuntimeError('401')

        sorted_map = OrderedDict(sorted(newResults.items(), reverse=True))
        ordered = list(sorted_map.values())

        return web.json_response({ "results": ordered }, headers={
            'Access-Control-Allow-Origin': '*'
        })

    except RuntimeError as e:
        error_message = e.args[0]
        if error_message == "401":
            return web.Response(text='AppleID 登录失效，请移除 auth.json 并重启 macless-haystack 服务', status=541)
    except Exception as e:
        logger.error(f"Unexpected error: {e}", exc_info=True)
        return web.Response(status=501)

def check_if_anisette_is_reachable(max_retries=3, retry_delay=10):
    server_url = mh_config.getAnisetteServer()
    logging.info(f'Checking if Anisette {server_url} is reachable')
    for attempt in range(max_retries):
        try:
            response = requests.get(server_url, timeout=5)
            response.raise_for_status()
            return
        except (requests.RequestException, requests.HTTPError) as e:
            logger.error(f"Attempt {attempt + 1} failed: {str(e)}")
            if attempt < max_retries - 1:
                logger.error(f"Retrying in {retry_delay} seconds...")
                time.sleep(retry_delay)
    logger.error(f"Max retries reached. Program will exit. Make sure your Anisette is reachable and start again with 'docker start -ai macless-haystack'")
    sys.exit()

if __name__ == "__main__":
    check_if_anisette_is_reachable()
    logging.info(f'Searching for token at ' + mh_config.getConfigFile())
    if not os.path.exists(mh_config.getConfigFile()):
        logging.info(f'No auth-token found.')
        apple_cryptography.registerDevice()

    app = web.Application()
    app.router.add_options('/', handle_options)
    app.router.add_get('/', handle_get)
    app.router.add_post('/', handle_post)
    
    address = mh_config.getBindingAddress() + ":" + str(mh_config.getPort())
    hasCert = os.path.isfile(mh_config.getCertFile())
    if hasCert:
        logger.info("Certificate file " + mh_config.getCertFile() +
                    " exists, so using SSL")
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ssl_context.load_cert_chain(certfile=mh_config.getCertFile(
        ), keyfile=mh_config.getKeyFile() if os.path.isfile(mh_config.getKeyFile()) else None)

        logger.info("serving at " + address + " over HTTPS")
    else:
        logger.info("Certificate file " + mh_config.getCertFile() +
                    " not found, so not using SSL")
        logger.info("serving at " + address + " over HTTP")
    
    user = mh_config.getEndpointUser()
    passw = mh_config.getEndpointPass()
    if (user is None or user == "") and (passw is None or passw == ""):
        logger.warning("Endpoint is not protected by authentication")
    else:
        logger.info("Endpoint is protected by authentication")

    access_logger = logging.getLogger('aiohttp.access')
    access_logger.setLevel(logging.WARNING)

    app['sem'] = asyncio.Semaphore(4)  # 最多4个并发请求，防止被封IP
    web.run_app(app, host=mh_config.getBindingAddress(), port=mh_config.getPort(), ssl_context=ssl_context if hasCert else None)

